import cv2
import sys
import numpy as np
import time
from PyQt5.QtWidgets import QApplication
from PyQt5.QtGui import *
import win32gui
import sys
from screen import get_screen
from keys import key_check
import random

#actions table 
w = [1,0,0,0,0,0,0]
s = [0,1,0,0,0,0,0]
d = [0,0,1,0,0,0,0]
a = [0,0,0,1,0,0,0]
skill = [0,0,0,0,1,0,0]
ak = [0,0,0,0,0,1,0]
big_skill = [0,0,0,0,0,0,1]
data = []

def main():
    global data
    start_time = 5
    for i in range(start_time-1):
        print(5-i)
        time.sleep(1)
    getting_data = True
    while getting_data:
        a = time.time()
        screen = get_screen()
        key_list = key_check()
        action = key_to_table(key_list)
        data.append([screen,action])
        spend_time = time.time() - a
        print('spend time:',spend_time)
        if len(data)%500 == 0 :
            np.save('train_data_v1',data)
            print('save')
        
def key_to_table(key_list):
    if 'W' in key_list:
        return w
    elif 'S' in key_list:
        return s
    elif 'D' in key_list:
        return d
    elif 'A' in key_list:
        return a    
    elif 'J' in key_list:
        return skill
    elif 'K' in key_list:
        return ak
    elif 'L' in key_list:
        return big_skill
    else:
        return d
main()